{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 07: understanding restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "import matplotlib.collections\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.tri as tri\n",
    "import numpy as np\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from mpi4py import MPI\n",
    "from dolfinx import FunctionSpace\n",
    "from dolfinx.fem import locate_dofs_geometrical, locate_dofs_topological\n",
    "from dolfinx.io import XDMFFile\n",
    "from multiphenicsx.fem import DofMapRestriction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary functions for plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_mesh(mesh, ax=None):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    ax.set_aspect(\"equal\")\n",
    "    points = mesh.geometry.x\n",
    "    cells = mesh.geometry.dofmap.array.reshape((-1, mesh.topology.dim + 1))\n",
    "    tria = tri.Triangulation(points[:, 0], points[:, 1], cells)\n",
    "    ax.triplot(tria, color=\"k\")\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_mesh_tags(mesh_tags, ax=None):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    ax.set_aspect(\"equal\")\n",
    "    mesh = mesh_tags.mesh\n",
    "    points = mesh.geometry.x\n",
    "    colors = [\"b\", \"r\"]\n",
    "    cmap = matplotlib.colors.ListedColormap(colors)\n",
    "    cmap_bounds = [0, 0.5, 1]\n",
    "    norm = matplotlib.colors.BoundaryNorm(cmap_bounds, cmap.N)\n",
    "    assert mesh_tags.dim in (mesh.topology.dim, mesh.topology.dim - 1)\n",
    "    if mesh_tags.dim == mesh.topology.dim:\n",
    "        cells = mesh.geometry.dofmap.array.reshape((-1, mesh.topology.dim + 1))\n",
    "        tria = tri.Triangulation(points[:, 0], points[:, 1], cells)\n",
    "        cell_colors = np.zeros((cells.shape[0], ))\n",
    "        cell_colors[mesh_tags.indices[mesh_tags.values == 1]] = 1\n",
    "        mappable = ax.tripcolor(tria, cell_colors, edgecolor=\"k\", cmap=cmap, norm=norm)\n",
    "    elif mesh_tags.dim == mesh.topology.dim - 1:\n",
    "        tdim = mesh.topology.dim\n",
    "        geometry_dofmap = mesh.geometry.dofmap\n",
    "        cells_map = mesh.topology.index_map(mesh.topology.dim)\n",
    "        num_cells = cells_map.size_local + cells_map.num_ghosts\n",
    "        connectivity_cells_to_facets = mesh.topology.connectivity(tdim, tdim - 1)\n",
    "        connectivity_cells_to_vertices = mesh.topology.connectivity(tdim, 0)\n",
    "        connectivity_facets_to_vertices = mesh.topology.connectivity(tdim - 1, 0)\n",
    "        vertex_map = {topology_index: geometry_index\n",
    "                      for c in range(num_cells)\n",
    "                      for (topology_index, geometry_index) in zip(\n",
    "                          connectivity_cells_to_vertices.links(c), geometry_dofmap.links(c))}\n",
    "        linestyles = [(0, (5, 10)), \"solid\"]\n",
    "        lines = list()\n",
    "        lines_colors_as_int = list()\n",
    "        lines_colors_as_str = list()\n",
    "        lines_linestyles = list()\n",
    "        mesh_tags_1 = mesh_tags.indices[mesh_tags.values == 1]\n",
    "        for c in range(num_cells):\n",
    "            facets = connectivity_cells_to_facets.links(c)\n",
    "            for f in facets:\n",
    "                if f in mesh_tags_1:\n",
    "                    value_f = 1\n",
    "                else:\n",
    "                    value_f = 0\n",
    "                vertices = [vertex_map[v] for v in connectivity_facets_to_vertices.links(f)]\n",
    "                lines.append(points[vertices][:, :2])\n",
    "                lines_colors_as_int.append(value_f)\n",
    "                lines_colors_as_str.append(colors[value_f])\n",
    "                lines_linestyles.append(linestyles[value_f])\n",
    "        mappable = matplotlib.collections.LineCollection(lines, cmap=cmap, norm=norm,\n",
    "                                                         colors=lines_colors_as_str,\n",
    "                                                         linestyles=lines_linestyles)\n",
    "        mappable.set_array(np.array(lines_colors_as_int))\n",
    "        ax.add_collection(mappable)\n",
    "        ax.autoscale()\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "    plt.colorbar(mappable, cax=cax, boundaries=cmap_bounds, ticks=cmap_bounds)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _plot_dofmap(coordinates, ax=None):\n",
    "    if ax is None:\n",
    "        ax = plt.gca()\n",
    "    text_offset = [1e-2, 1e-2]\n",
    "    ax.scatter(coordinates[:, 0], coordinates[:, 1], c=\"k\", s=50)\n",
    "    for c in np.unique(coordinates, axis=0):\n",
    "        dofs_c = (coordinates == c).all(axis=1).nonzero()[0]\n",
    "        text_c = np.array2string(dofs_c, separator=\", \", max_line_width=10)\n",
    "        ax.text(c[0] + text_offset[0], c[1] + text_offset[1], text_c, fontsize=20)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_dofmap(V, ax=None):\n",
    "    coordinates = V.tabulate_dof_coordinates().round(decimals=3)\n",
    "    return _plot_dofmap(coordinates, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_dofmap_restriction(V, restriction, ax=None):\n",
    "    coordinates = V.tabulate_dof_coordinates().round(decimals=3)\n",
    "    return _plot_dofmap(coordinates[list(restriction.unrestricted_to_restricted.keys())], ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary functions for asserts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_dofs(restriction, comm):\n",
    "    u2r = restriction.unrestricted_to_restricted\n",
    "    restricted_global_indices = restriction.index_map.global_indices()\n",
    "    dofs_V_restriction_global = [restricted_global_indices[dof] for (_, dof) in u2r.items()]\n",
    "    return len(set(gdof\n",
    "                   for gdofs in comm.allgather(dofs_V_restriction_global)\n",
    "                   for gdof in gdofs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def locate_dofs_by_polar_coordinates(r, theta, V, restriction):\n",
    "    p = [r * np.cos(theta), r * np.sin(theta), 0.]\n",
    "    dofs_V = locate_dofs_geometrical(V, lambda x: np.isclose(x.T, p).all(axis=1)).reshape(-1)\n",
    "    u2r = restriction.unrestricted_to_restricted\n",
    "    restricted_global_indices = restriction.index_map.global_indices()\n",
    "    dofs_V_restriction = [u2r[dof] for dof in dofs_V if dof in u2r]\n",
    "    dofs_V_restriction_global = [restricted_global_indices[dof] for dof in dofs_V_restriction]\n",
    "    return set(gdof\n",
    "               for gdofs in V.mesh.mpi_comm().allgather(dofs_V_restriction_global)\n",
    "               for gdof in gdofs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh\n",
    "\n",
    "Read in a simple mesh consisting in an hexagon discretized with six equilateral triangle cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/hexagon.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh()\n",
    "mesh.topology.create_connectivity_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh restrictions on cells\n",
    "\n",
    "Read in mesh tags on cells, which are equal to one on all cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/hexagon.xdmf\", \"r\") as infile:\n",
    "    cell_restriction_all = infile.read_meshtags(mesh, name=\"cell_restriction_all\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mesh_tags(cell_restriction_all)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read in mesh tags on cells, which are equal to one on one half of the cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/hexagon.xdmf\", \"r\") as infile:\n",
    "    cell_restriction_subset = infile.read_meshtags(mesh, name=\"cell_restriction_subset\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mesh_tags(cell_restriction_subset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh restrictions on facets\n",
    "\n",
    "Read in mesh tags on facets, which are equal to one on all facets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/hexagon.xdmf\", \"r\") as infile:\n",
    "    facet_restriction_all = infile.read_meshtags(mesh, name=\"facet_restriction_all\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mesh_tags(facet_restriction_all)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read in mesh tags on facets, which are equal to one on two facets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/hexagon.xdmf\", \"r\") as infile:\n",
    "    facet_restriction_subset = infile.read_meshtags(mesh, name=\"facet_restriction_subset\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_mesh_tags(facet_restriction_subset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_restrictions = (cell_restriction_all, cell_restriction_subset)\n",
    "facet_restrictions = (facet_restriction_all, facet_restriction_subset)\n",
    "all_restrictions = cell_restrictions + facet_restrictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Lagrange spaces\n",
    "\n",
    "Define Lagrange FE spaces of order $k=1, 2, 3$, and plot the associated DofMap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG = [FunctionSpace(mesh, (\"Lagrange\", k)) for k in (1, 2, 3)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 1, figsize=(10, 30))\n",
    "for (k, CGk) in enumerate(CG):\n",
    "    plot_mesh(mesh, ax[k])\n",
    "    plot_dofmap(CGk, ax[k])\n",
    "    ax[k].set_title(\"CG \" + str(k + 1) + \" DofMap\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define DofMapRestriction objects associated to the Lagrange FE spaces, for all four restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dofmap_restriction_CG = dict()\n",
    "for restriction in all_restrictions:\n",
    "    dofmap_restriction_CG[restriction] = list()\n",
    "    for CGk in CG:\n",
    "        restrict_CGk = locate_dofs_topological(\n",
    "            CGk, restriction.dim, restriction.indices[restriction.values == 1])\n",
    "        dofmap_restriction_CG[restriction].append(DofMapRestriction(CGk.dofmap, restrict_CGk))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for the case of cell restriction equal to one on the entire domain. There is indeed no difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, CGk) in enumerate(CG):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(CGk, ax[k, 0])\n",
    "    ax[k, 0].set_title(\"CG \" + str(k + 1) + \" DofMap\", fontsize=30)\n",
    "for (k, (CGk, dofmap_restriction_CGk)) in enumerate(zip(CG, dofmap_restriction_CG[cell_restriction_all])):\n",
    "    plot_mesh_tags(cell_restriction_all, ax[k, 1])\n",
    "    plot_dofmap_restriction(CGk, dofmap_restriction_CGk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"CG \" + str(k + 1) + \" DofMapRestriction\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_1 = CG[0]\n",
    "dofmap_restriction_CG_1 = dofmap_restriction_CG[cell_restriction_all][0]\n",
    "assert count_dofs(dofmap_restriction_CG_1, mesh.mpi_comm()) == 7\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_2 = CG[1]\n",
    "dofmap_restriction_CG_2 = dofmap_restriction_CG[cell_restriction_all][1]\n",
    "assert count_dofs(dofmap_restriction_CG_2, mesh.mpi_comm()) == 19\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 3 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 11 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_3 = CG[2]\n",
    "dofmap_restriction_CG_3 = dofmap_restriction_CG[cell_restriction_all][2]\n",
    "assert count_dofs(dofmap_restriction_CG_3, mesh.mpi_comm()) == 37\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 3 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 5 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 7 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 9 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 11 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 2 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 2 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 4 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 4 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 5 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 5 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for che case of cell restriction equal to one on a subset of the domain. Note how the DofMapRestriction has only a subset of the DOFs of the DofMap, and properly renumbers them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, CGk) in enumerate(CG):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(CGk, ax[k, 0])\n",
    "    ax[k, 0].set_title(\"CG \" + str(k + 1) + \" DofMap\", fontsize=30)\n",
    "for (k, (CGk, dofmap_restriction_CGk)) in enumerate(zip(CG, dofmap_restriction_CG[cell_restriction_subset])):\n",
    "    plot_mesh_tags(cell_restriction_subset, ax[k, 1])\n",
    "    plot_dofmap_restriction(CGk, dofmap_restriction_CGk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"CG \" + str(k + 1) + \" DofMapRestriction\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_1 = CG[0]\n",
    "dofmap_restriction_CG_1 = dofmap_restriction_CG[cell_restriction_subset][0]\n",
    "assert count_dofs(dofmap_restriction_CG_1, mesh.mpi_comm()) == 5\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_2 = CG[1]\n",
    "dofmap_restriction_CG_2 = dofmap_restriction_CG[cell_restriction_subset][1]\n",
    "assert count_dofs(dofmap_restriction_CG_2, mesh.mpi_comm()) == 12\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_3 = CG[2]\n",
    "dofmap_restriction_CG_3 = dofmap_restriction_CG[cell_restriction_subset][2]\n",
    "assert count_dofs(dofmap_restriction_CG_3, mesh.mpi_comm()) == 22\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 5 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 7 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 9 * np.pi / 6,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 2 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 4 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 4 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 5 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for che case of facet restriction equal to one on the entire domain. Note how there is no difference for $k=1, 2$, but the cases $k=3$ differ (the DofMapRestriction does not have the DOF at the cell center)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, CGk) in enumerate(CG):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(CG[k], ax[k, 0])\n",
    "    ax[k, 0].set_title(\"CG \" + str(k + 1) + \" DofMap\", fontsize=30)\n",
    "for (k, (CGk, dofmap_restriction_CGk)) in enumerate(zip(CG, dofmap_restriction_CG[facet_restriction_all])):\n",
    "    plot_mesh_tags(facet_restriction_all, ax[k, 1])\n",
    "    plot_dofmap_restriction(CGk, dofmap_restriction_CGk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"CG \" + str(k + 1) + \" DofMapRestriction\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_1 = CG[0]\n",
    "dofmap_restriction_CG_1 = dofmap_restriction_CG[facet_restriction_all][0]\n",
    "assert count_dofs(dofmap_restriction_CG_1, mesh.mpi_comm()) == 7\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_2 = CG[1]\n",
    "dofmap_restriction_CG_2 = dofmap_restriction_CG[facet_restriction_all][1]\n",
    "assert count_dofs(dofmap_restriction_CG_2, mesh.mpi_comm()) == 19\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 3 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 11 * np.pi / 6,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_3 = CG[2]\n",
    "dofmap_restriction_CG_3 = dofmap_restriction_CG[facet_restriction_all][2]\n",
    "assert count_dofs(dofmap_restriction_CG_3, mesh.mpi_comm()) == 31\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 2 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 2 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 4 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 4 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 5 * np.pi / 3 - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, 5 * np.pi / 3 + np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(7) / 3, np.pi - np.arctan(np.sqrt(3) / 5),\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for che case of facet restriction equal to one on a subset of the domain. Note how the DofMapRestriction has only a subset of the DOFs of the DofMap, and properly renumbers them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, CGk) in enumerate(CG):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(CG[k], ax[k, 0])\n",
    "    ax[k, 0].set_title(\"CG \" + str(k + 1) + \" DofMap\", fontsize=30)\n",
    "for (k, (CGk, dofmap_restriction_CGk)) in enumerate(zip(CG, dofmap_restriction_CG[facet_restriction_subset])):\n",
    "    plot_mesh_tags(facet_restriction_subset, ax[k, 1])\n",
    "    plot_dofmap_restriction(CGk, dofmap_restriction_CGk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"CG \" + str(k + 1) + \" DofMapRestriction\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_1 = CG[0]\n",
    "dofmap_restriction_CG_1 = dofmap_restriction_CG[facet_restriction_subset][0]\n",
    "assert count_dofs(dofmap_restriction_CG_1, mesh.mpi_comm()) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_1, dofmap_restriction_CG_1)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_2 = CG[1]\n",
    "dofmap_restriction_CG_2 = dofmap_restriction_CG[facet_restriction_subset][1]\n",
    "assert count_dofs(dofmap_restriction_CG_2, mesh.mpi_comm()) == 5\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_2, dofmap_restriction_CG_2)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CG_3 = CG[2]\n",
    "dofmap_restriction_CG_3 = dofmap_restriction_CG[facet_restriction_subset][2]\n",
    "assert count_dofs(dofmap_restriction_CG_3, mesh.mpi_comm()) == 7\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(2 / 3, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            CG_3, dofmap_restriction_CG_3)) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discontinuous Galerkin spaces\n",
    "\n",
    "Define Discontinuous Galerkin spaces of order $k = 0, 1, 2$, and plot the associated DofMap. This spaces will be used in combination with a cell restriction, as DG DOFs are associated to cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG = [FunctionSpace(mesh, (\"Discontinuous Lagrange\", k)) for k in (0, 1, 2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 1, figsize=(10, 30))\n",
    "for (k, DGk) in enumerate(DG):\n",
    "    plot_mesh(mesh, ax[k])\n",
    "    plot_dofmap(DGk, ax[k])\n",
    "    ax[k].set_title(\"DG \" + str(k) + \" DofMap\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define Discontinuous Galerkin Trace spaces of order $k = 0, 1, 2, 3$, and plot the associated DofMap. This spaces will be used in combination with a facet restriction, as DGT DOFs are associated to facets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT = [FunctionSpace(mesh, (\"Discontinuous Lagrange Trace\", k)) for k in (0, 1, 2)]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "_, ax = plt.subplots(3, 1, figsize=(10, 30))\n",
    "for (k, DGTk) in enumerate(DGT):\n",
    "    plot_mesh(mesh, ax[k])\n",
    "    plot_dofmap(DGTk, ax[k])\n",
    "    ax[k].set_title(\"DGT \" + str(k) + \" DofMap\\n\", fontsize=30)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define DofMapRestriction objects associated to the Discontinuos Galerkin FE spaces, for all cell restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dofmap_restriction_DG = dict()\n",
    "for restriction in cell_restrictions:\n",
    "    dofmap_restriction_DG[restriction] = list()\n",
    "    for DGk in DG:\n",
    "        restrict_DGk = locate_dofs_topological(\n",
    "            DGk, restriction.dim, restriction.indices[restriction.values == 1])\n",
    "        dofmap_restriction_DG[restriction].append(DofMapRestriction(DGk.dofmap, restrict_DGk))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define DofMapRestriction objects associated to the Discontinuos Galerkin Trace FE spaces, for all facet restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "dofmap_restriction_DGT = dict()\n",
    "for restriction in facet_restrictions:\n",
    "    dofmap_restriction_DGT[restriction] = list()\n",
    "    for DGTk in DGT:\n",
    "        restrict_DGTk = locate_dofs_topological(\n",
    "            DGTk, restriction.dim, restriction.indices[restriction.values == 1])\n",
    "        dofmap_restriction_DGT[restriction].append(DofMapRestriction(DGTk.dofmap, restrict_DGTk))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for the case of cell restriction equal to one on the entire domain. There is indeed no difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, DGk) in enumerate(DG):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(DGk, ax[k, 0])\n",
    "    ax[k, 0].set_title(\"DG \" + str(k) + \" DofMap\", fontsize=30)\n",
    "for (k, (DGk, dofmap_restriction_DGk)) in enumerate(zip(DG, dofmap_restriction_DG[cell_restriction_all])):\n",
    "    plot_mesh_tags(cell_restriction_all, ax[k, 1])\n",
    "    plot_dofmap_restriction(DGk, dofmap_restriction_DGk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"DG \" + str(k) + \" DofMapRestriction\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG_0 = DG[0]\n",
    "dofmap_restriction_DG_0 = dofmap_restriction_DG[cell_restriction_all][0]\n",
    "assert count_dofs(dofmap_restriction_DG_0, mesh.mpi_comm()) == 6\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 3 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 5 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 7 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 9 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 11 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG_1 = DG[1]\n",
    "dofmap_restriction_DG_1 = dofmap_restriction_DG[cell_restriction_all][1]\n",
    "assert count_dofs(dofmap_restriction_DG_1, mesh.mpi_comm()) == 18\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 6\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG_2 = DG[2]\n",
    "dofmap_restriction_DG_2 = dofmap_restriction_DG[cell_restriction_all][2]\n",
    "assert count_dofs(dofmap_restriction_DG_2, mesh.mpi_comm()) == 36\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 6\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 0,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 3 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 11 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for che case of cell restriction equal to one on a subset of the domain. Note how the DofMapRestriction has only a subset of the DOFs of the DofMap, and properly renumbers them. Note also that the number of DOFs at the same physical location might be different between DofMap and DofMapRestriction (see e.g. the center of the hexagon)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, DGk) in enumerate(DG):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(DGk, ax[k, 0])\n",
    "    ax[k, 0].set_title(\"DG \" + str(k) + \" DofMap\", fontsize=30)\n",
    "for (k, (DGk, dofmap_restriction_DGk)) in enumerate(zip(DG, dofmap_restriction_DG[cell_restriction_subset])):\n",
    "    plot_mesh_tags(cell_restriction_subset, ax[k, 1])\n",
    "    plot_dofmap_restriction(DGk, dofmap_restriction_DGk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"DG \" + str(k) + \" DofMapRestriction\", fontsize=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG_0 = DG[0]\n",
    "dofmap_restriction_DG_0 = dofmap_restriction_DG[cell_restriction_subset][0]\n",
    "assert count_dofs(dofmap_restriction_DG_0, mesh.mpi_comm()) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 5 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 7 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 3, 9 * np.pi / 6,\n",
    "                                            DG_0, dofmap_restriction_DG_0)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG_1 = DG[1]\n",
    "dofmap_restriction_DG_1 = dofmap_restriction_DG[cell_restriction_subset][1]\n",
    "assert count_dofs(dofmap_restriction_DG_1, mesh.mpi_comm()) == 9\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DG_1, dofmap_restriction_DG_1)) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DG_2 = DG[2]\n",
    "dofmap_restriction_DG_2 = dofmap_restriction_DG[cell_restriction_subset][2]\n",
    "assert count_dofs(dofmap_restriction_DG_2, mesh.mpi_comm()) == 18\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            DG_2, dofmap_restriction_DG_2)) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for che case of facet restriction equal to one on the entire domain. There is indeed no difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, DGTk) in enumerate(DGT):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(DGTk, ax[k, 0])\n",
    "    ax[k, 0].set_title(\"DGT \" + str(k) + \" DofMap\\n\", fontsize=30)\n",
    "for (k, (DGTk, dofmap_restriction_DGTk)) in enumerate(zip(DGT, dofmap_restriction_DGT[facet_restriction_all])):\n",
    "    plot_mesh_tags(facet_restriction_all, ax[k, 1])\n",
    "    plot_dofmap_restriction(DGTk, dofmap_restriction_DGTk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"DGT \" + str(k) + \" DofMapRestriction\\n\", fontsize=30)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT_0 = DGT[0]\n",
    "dofmap_restriction_DGT_0 = dofmap_restriction_DGT[facet_restriction_all][0]\n",
    "assert count_dofs(dofmap_restriction_DGT_0, mesh.mpi_comm()) == 12\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 0,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi / 3,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, np.pi / 6,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 3 * np.pi / 6,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 11 * np.pi / 6,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT_1 = DGT[1]\n",
    "dofmap_restriction_DGT_1 = dofmap_restriction_DGT[facet_restriction_all][1]\n",
    "assert count_dofs(dofmap_restriction_DGT_1, mesh.mpi_comm()) == 24\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 6\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 3\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT_2 = DGT[2]\n",
    "dofmap_restriction_DGT_2 = dofmap_restriction_DGT[facet_restriction_all][2]\n",
    "assert count_dofs(dofmap_restriction_DGT_2, mesh.mpi_comm()) == 36\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 6\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 0,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, np.pi,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 4 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 0,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, np.pi,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 4 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 3\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, np.pi / 6,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 3 * np.pi / 6,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 5 * np.pi / 6,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 7 * np.pi / 6,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 9 * np.pi / 6,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(np.sqrt(3) / 2, 11 * np.pi / 6,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare DOFs for che case of facet restriction equal to one on a subset of the domain. Note how the DofMapRestriction has only a subset of the DOFs of the DofMap, and properly renumbers them. Note also that the number of DOFs at the same physical location might be different between DofMap and DofMapRestriction (see e.g. the center of the hexagon)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "_, ax = plt.subplots(3, 2, figsize=(20, 30))\n",
    "for (k, DGTk) in enumerate(DGT):\n",
    "    plot_mesh(mesh, ax[k, 0])\n",
    "    plot_dofmap(DGTk, ax[k, 0])\n",
    "    ax[k, 0].set_title(\"DGT \" + str(k) + \" DofMap\\n\", fontsize=30)\n",
    "for (k, (DGTk, dofmap_restriction_DGTk)) in enumerate(zip(DGT, dofmap_restriction_DGT[facet_restriction_subset])):\n",
    "    plot_mesh_tags(facet_restriction_subset, ax[k, 1])\n",
    "    plot_dofmap_restriction(DGTk, dofmap_restriction_DGTk, ax[k, 1])\n",
    "    ax[k, 1].set_title(\"DGT \" + str(k) + \" DofMapRestriction\\n\", fontsize=30)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Assert that DOFs are at the expected locations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT_0 = DGT[0]\n",
    "dofmap_restriction_DGT_0 = dofmap_restriction_DGT[facet_restriction_subset][0]\n",
    "assert count_dofs(dofmap_restriction_DGT_0, mesh.mpi_comm()) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            DGT_0, dofmap_restriction_DGT_0)) == 1\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT_1 = DGT[1]\n",
    "dofmap_restriction_DGT_1 = dofmap_restriction_DGT[facet_restriction_subset][1]\n",
    "assert count_dofs(dofmap_restriction_DGT_1, mesh.mpi_comm()) == 4\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DGT_1, dofmap_restriction_DGT_1)) == 1\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" # Temporarily disabled due to lack of suport for DGT in basix\n",
    "DGT_2 = DGT[2]\n",
    "dofmap_restriction_DGT_2 = dofmap_restriction_DGT[facet_restriction_subset][2]\n",
    "assert count_dofs(dofmap_restriction_DGT_2, mesh.mpi_comm()) == 6\n",
    "assert len(locate_dofs_by_polar_coordinates(0, 0,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 2\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 2 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(0.5, 5 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 2 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "assert len(locate_dofs_by_polar_coordinates(1, 5 * np.pi / 3,\n",
    "                                            DGT_2, dofmap_restriction_DGT_2)) == 1\n",
    "\"\"\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
